#include <assert.h>
#include <test/testCases.cuh>
#include <iostream>
#include <utils/Profiler.h>
#include <test/testMyData.cuh>
#include <DataConvert.cuh>
#include <utils/DataCatcher.h>
#include <Precond/RugeStuben.cuh>
#include <Solver/BiCGStabSolver.cuh>
#include <Data/MatrixFeatures.cuh>

#include "CSRPrecondAMG.cuh"

#define CHECK(call)                                                               \
{                                                                                 \
const cudaError_t error = call;                                               \
if (error != cudaSuccess)                                                     \
{                                                                             \
printf("Error: %s:%d, ", __FILE__, __LINE__);                             \
printf("code: %d, reason: %s\n", error, cudaGetErrorString(error));       \
exit(1);                                                                  \
}                                                                             \
}


template<typename T>
__global__ void printKernel(T *data, const size_t N) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        printf("%f\n", static_cast<double>(data[idx]));
    }
}


int main(int argc, char *argv[]) {
    /*
    MPI_Init(&argc, &argv);
    profileJacobiBiCGStab(argc, argv);
    profileSORBiCGStab(argc, argv);
    auto A = myamg::CSR2Matrix<int, double, myamg::MatrixDeviceAllocator>("Vely-1.csr");
    auto b = myamg::CSRRhs2Vector<double, myamg::MatrixDeviceAllocator>("Vely-1.csr");
    myamg::Vector<double, myamg::MatrixDeviceAllocator> x(A.numRows());
    myamg::vectorClear(x);
    std::cout << "Matrix loaded." << std::endl;
    std::cout << "Rows: " << A.numRows() << ", Cols: " << A.numCols() << ", NNZ: " << A.numNonzeros() << std::endl;
    std::cout << "Rhs Rows:" << b.size() << std::endl;
    myamg::CSRPrecondAMG<int, double, myamg::MatrixDeviceAllocator> amg{};
    amg.coarsenType = 0;
    amg.interpType = 0;
    amg.aggressiveLevels = 50;
    amg.interpMaxElements = 4;
    amg.interpMinElements = 4;
    amg.truncationFactor = 0.1;
    amg.printStats = true;
    {
        PROFILE_SCOPE("Jacobi-AMG Setup GPU");
        amg.setup(A);
    }
    std::vector<std::vector<std::pair<std::string, std::string>> > features;
    for (auto i = 0; i < amg.getLevels(); i++) {
        const auto& A = amg.getA(i);
        myamg::MatrixFeatures mf{};
        mf.compute(A);
        features.push_back(mf.toMap());
    }
    myamg::writeFeaturesToCSV(features, "features.csv");
    std::cout << "AMG setup completed." << std::endl;
    std::cout << "Solve started." << std::endl;
    myamg::BiCGStabSolver<int, double, myamg::MatrixDeviceAllocator> solver{};
    int iterations = 0;
    double relres = 0.0;
    solver.printStats = true;
    {
        PROFILE_SCOPE("Jacobi-AMG-BiCGStab Solve GPU");
        solver.apply(A, amg, b, x, iterations, relres);
    }
    myamg::Profiler::instance().saveToFile("Result/Solve_Multi.txt", false);
    MPI_Finalize();
    */
    auto matrixHost = myamg::Mtx2Matrix<int, double>("A_1.mtx");
    auto vectorHost = myamg::Mtx2Vector<double>("b_1.mtx");
    auto matrixDevice = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("A_1.mtx");
    auto vectorDevice = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("b_1.mtx");
    std::cout << "Matrix loaded." << std::endl;
    const auto numRows = matrixHost.numRows();
    auto Ap = matrixHost.rowPtr();
    auto Ai = matrixHost.colIndex();
    auto Av = matrixHost.values();
    myamg::Vector<double> result1(numRows);
    myamg::Vector<double> result2(numRows);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> result3(numRows);
    myamg::TensorSpMVHandler<myamg::MatrixHostAllocator> tensorHandler;
    std::cout << "Tensor SpMV Setup starts." << std::endl;
    tensorHandler.setup(matrixHost);
    std::cout << "Tensor SpMV Setup completed." << std::endl;
    for (auto i = 0; i < 10; ++i) {
        myamg::matMulVec(matrixHost, vectorHost, result1);
        myamg::matMulVec(matrixDevice, vectorDevice, result3);
        tensorHandler.apply(vectorHost, result2);
    }
    // print and compare results
    unsigned int numShortRows = 0;
    unsigned int numLongRows = 0;
    unsigned int numVeryLongRows = 0;
    for (unsigned int i = 0; i < numRows; ++i) {
        const auto numNoneZero = Ap[i + 1] - Ap[i];
        if (numNoneZero > 0) {
            if (numNoneZero <= 8) {
                ++numShortRows;
            } else if (numNoneZero <= 12) {
                ++numLongRows;
            } else {
                ++numVeryLongRows;
            }
        }
    }
    // 建立映射
    auto* shortRowMap = new unsigned int[numShortRows];
    auto* longRowMap = new unsigned int[numLongRows];
    auto* veryLongRowMap = new unsigned int[numVeryLongRows];
    unsigned int shortIdx = 0;
    unsigned int longIdx = 0;
    unsigned int veryLongIdx = 0;
    for (unsigned int i = 0; i < numRows; ++i) {
        const auto numNoneZero = Ap[i + 1] - Ap[i];
        if (numNoneZero > 0) {
            if (numNoneZero <= 8) {
                shortRowMap[shortIdx++] = i;
            } else if (numNoneZero <= 12) {
                longRowMap[longIdx++] = i;
            } else {
                veryLongRowMap[veryLongIdx++] = i;
            }
        }
    }
    double eps = 1e-6;
    for (auto i = 0; i < numLongRows; ++i) {
        const auto realRow = longRowMap[i];
        assert(std::abs(result1.values()[realRow] - result2.values()[realRow]) < eps);
    }
    std::cout << "Long rows check passed." << std::endl;
    for (auto i = 0; i < numVeryLongRows; ++i) {
        const auto realRow = veryLongRowMap[i];
        assert(std::abs(result1.values()[realRow] - result2.values()[realRow]) < eps);
    }
    std::cout << "Very long rows check passed." << std::endl;
    for (auto i = 0; i < numShortRows; ++i) {
        const auto realRow = shortRowMap[i];
        assert(std::abs(result1.values()[realRow] - result2.values()[realRow]) < eps);
    }
    std::cout << "Short rows check passed." << std::endl;
    std::cout << "All rows check passed." << std::endl;
    return 0;
}
